#####################################################
#
# Python2 script to reproduce Fig. 2 in the comment
# authors: Thomas Lang, Stephan Hesselmann
#
#####################################################
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import glob

ecolor = ['#cc2909', '#efc600', '#60af92', '#3e77af','#3e77af','#60af92','#efc600','#cc2909']
fcolor = ['#ea6868', '#eddea2', '#99d1b9', '#a3c1e5','#a3c1e5','#99d1b9','#eddea2','#ea6868']
marker = ['o','s','D','^','o','s','D','^']
markerSize = np.array([5,4.5,4,5.5,5,4.5,4,5.5])*2.0

fsTicksLabel = 18
fsTicksLabelInset = 16
fsAxesLabel = 21
fsLegend = 16
fsLabel = 18

a1 = np.array([1./2., -np.sqrt(3.)/2.])
a2 = np.array([1., 0.])
b1 = np.array([2.*np.pi, -2.*np.pi/3.*np.sqrt(3.)])
b2 = np.array([0., 4.*np.pi/np.sqrt(3.)])
KK = (2*b2 + b1)/3.
K = { 6 : np.array([4.1887902, 0.]), 9 : np.array([4.1887902, 0.]), 12 : np.array([4.1887902, 0.]), 15 : np.array([4.1887902, 0.]) }
def E(k):
    ek = np.zeros(len(k))
    for ik in range(len(k)):
        ek[ik] = np.abs(1. + np.exp(1.j * np.dot(k[ik], -a2)) + np.exp(1.j * np.dot(k[ik], a1 - a2)))
    return ek

i_L = 0
i_gamma = 1
i_U = 2
i_q = 3
i_kx = 4
i_ky = 5
i_delta = 6
i_sigma = 7

L_list = [ 15 ]
gamma_list = [ 0. ]
U_list = [ 0.5, 1., 1.5, 2., 2.5, 3., 3.4, 3.6, 3.75, 4, 4.5 ]
q_list = [ 0, 1, 2]

fig = plt.figure(figsize=(18*0.7,8*0.7), dpi=96)

#------------------------------------------------------------------------------------

ax = plt.subplot(131)
ax.tick_params(which='both', axis='both', direction='in', bottom=True, top=True, left=True, right=True, labelsize=fsTicksLabel)
ax.set_xlabel(r'$U/t$', fontsize=fsAxesLabel)
ax.set_ylabel(r'$E/t$', fontsize=fsAxesLabel)
ax.set_yticks([0, 0.4, 0.8, 1.2])
ax.yaxis.set_minor_locator(ticker.AutoMinorLocator(4))
ax.xaxis.set_minor_locator(ticker.AutoMinorLocator(2))
ax.set_xlim(0.3, 4.7)
ax.set_ylim(-0.05, 1.2)
props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
ax.text(-0.26, 0.98, r'$\mathbf{A}$', color='#000000', transform=ax.transAxes, fontSize=21, ha='left')
ax.text( 1.10, 0.98, r'$\mathbf{B}$', color='#000000', transform=ax.transAxes, fontSize=21, ha='left')
ax.text( 3.0, 0.98, r'$\mathbf{C}$', color='#000000', transform=ax.transAxes, fontSize=21, ha='left')
ax.text(0.08, 0.91, r'$\gamma = 0, L=15$', color='#000000', transform=ax.transAxes, fontSize=fsLabel, ha='left', bbox=dict(boxstyle='round', ec=(219/255., 202/255., 175/255.), fc=(244/255., 232/255., 211/255.)))
ax.text(0.20, 0.75, r'$a\Delta k = 0.97$', transform=ax.transAxes, fontsize=16, verticalalignment='top')
ax.text(0.20, 0.44, r'$a\Delta k = 0.48$', transform=ax.transAxes, fontsize=16, verticalalignment='top')
ax.text(0.20, 0.12, r'$a\Delta k = 0$', transform=ax.transAxes, fontsize=16, verticalalignment='top')
ax.text(0.715, 0.985, r'$U_c/t = 3.85$', transform=ax.transAxes, fontsize=16, verticalalignment='top', color='#888888', rotation=90)

plt.axhline(0, color='#dddddd', zorder=-2, linewidth=0.5)
plt.axvline(3.85, linestyle='--', color='#888888', linewidth=1.0)

filelist = glob.glob('data.txt')
filelist.sort()

cnt_q = 3
for filename in filelist:
	with open(filename) as f:
		lines = (line for line in f if not line.startswith('L'))
		new_del = []
		for l in lines:
			new_del.append(','.join(l.split()))
		data = np.loadtxt(new_del, delimiter=',', skiprows=0)

	data_filter = list(filter(lambda x : x[i_L] in L_list and x[i_gamma] in gamma_list and x[i_U] in U_list, data))

	for L in L_list:
		data_filter_L = np.array(list(filter(lambda x : x[i_L] == L and x[i_U] in U_list, data_filter)))
		
		for q_i in q_list:
			data_filter_q = np.array(list(filter(lambda x : (np.abs(K[L]-b1/L*q_i - np.array([x[i_kx], x[i_ky]])) < 1E-6).all().any(), data_filter_L)))
			delta_list = data_filter_q[:, i_delta]
			sigma_list = data_filter_q[:, i_sigma]			
			ax.errorbar(U_list, delta_list, yerr=sigma_list, markeredgewidth=1.6, capsize=8, fmt=marker[cnt_q]+'-', markersize=markerSize[cnt_q], color=ecolor[cnt_q], ecolor=ecolor[cnt_q], mfc=fcolor[cnt_q])


#------------------------------------------------------------------------------------

L_list = [6, 9, 12, 15]
gamma_list = [ 0. ]
U_list = [ 0.5, 1., 1.5, 2., 2.5, 3., 3.4, 3.6, 3.75, 4 ]
q_list = [ 1 ]
v0 = 3.**0.5/2.

ax2 = plt.subplot(132)
ax2.tick_params(which='both', axis='both', direction='in', bottom=True, top=True, left=True, right=True, labelleft=True,  labelright=False, labelsize=fsTicksLabel)
ax2.set_xlabel(r'$U/t$', fontsize=fsAxesLabel)
ax2.set_ylabel(r'$\;\;\;\Delta v/v_0 = [E(a\Delta\mathbf{k}) / a\Delta\mathbf{k} - v_0] / v_0$', fontsize=fsLabel, labelpad=5)
ax2.yaxis.set_label_position('left')
ax2.set_yticks([-0.5, -0.4, -0.3, -0.2, -0.1, 0, 0.1, 0.2, 0.3])
ax2.yaxis.set_minor_locator(ticker.AutoMinorLocator(2))
ax2.xaxis.set_minor_locator(ticker.AutoMinorLocator(2))
ax2.set_xlim(0.3, 4.15)
ax2.set_ylim(-0.51, 0.11)

ax2.text(0.08, 0.065, r'$\gamma = 0$', color='#000000', transform=ax2.transAxes, fontSize=fsLabel, ha='left', bbox=dict(boxstyle='round', ec=(219/255., 202/255., 175/255.), fc=(244/255., 232/255., 211/255.)))
ax2.text(0.825, 0.585, r'$U_c/t = 3.85$', transform=ax2.transAxes, fontsize=16, verticalalignment='top', color='#888888', rotation=90)
ax2.text(0.31, 0.16, r'40% reduced $v_0$', color='#cc2909', transform=ax2.transAxes, fontSize=16, ha='left')
ax2.annotate("", xy=(3.8, -0.4), xytext=(3.3, -0.4), arrowprops=dict(arrowstyle="->", linewidth=1.2, color='#cc2909'))
plt.axhline(0, color='#dddddd', zorder=-2, linewidth=0.5)
plt.axvline(3.85, linestyle='--', color='#888888', linewidth=1.0)


data_filter = list(filter(lambda x : x[i_L] in L_list and x[i_gamma] in gamma_list and x[i_U] in U_list, data))

cnt_q = 0
for L in L_list:
	data_filter_L = np.array(list(filter(lambda x : x[i_L] == L and x[i_U] in U_list, data_filter)))
	
	for q_i in q_list:
		data_filter_q = np.array(list(filter(lambda x : (np.abs(K[L]-b1/L*q_i - np.array([x[i_kx], x[i_ky]])) < 1E-6).all().any(), data_filter_L)))
		data_filter_K = np.array(list(filter(lambda x : (np.abs(K[L]          - np.array([x[i_kx], x[i_ky]])) < 1E-6).all().any(), data_filter_L)))

		if len(data_filter_q > 0):
			offset = v0 * q_i
			delta_k = np.linalg.norm(b1) / L
			delta_list = (data_filter_q[:, i_delta] / delta_k - offset) / v0
			sigma_list = data_filter_q[:, i_sigma] / delta_k / v0			
			ax2.plot(U_list, delta_list, color=ecolor[cnt_q])
			ax2.errorbar(U_list, delta_list, yerr=sigma_list, markeredgewidth=1.6, capsize=8, fmt=marker[cnt_q], markersize=markerSize[cnt_q], color=ecolor[cnt_q], ecolor=ecolor[cnt_q], mfc=fcolor[cnt_q])

			cnt_q += 1

cnt_q = 3
for L in np.flip(L_list,0):
    ax2.errorbar([-1,-1,-1,-1], [0,0,0,0], yerr=[0,0,0,0], markeredgewidth=1.6, capsize=8, fmt=marker[cnt_q], markersize=markerSize[cnt_q], color=ecolor[cnt_q], ecolor=ecolor[cnt_q], mfc=fcolor[cnt_q], label=r'$L=%d$'%L)
    cnt_q -= 1

plt.legend(loc=(0.03, 0.26), prop={'size': fsLegend}, frameon=False, labelspacing=0.4, handletextpad=-0.15, fontsize=fsLegend)


#------------------------------------------------------------------------------------

ax3 = plt.subplot(133)
ax3.tick_params(which='both', axis='both', direction='in', bottom=True, top=True, left=True, right=True, labelleft=True,  labelright=False, labelsize=fsTicksLabel)
ax3.set_xlabel(r'$U/t$', fontsize=fsAxesLabel)
ax3.set_ylabel(r'$\Delta v/v_0\;\;\;[\rm see\; Tang\; {\it et\; al.}]$', fontsize=fsLabel, labelpad=5)
ax3.yaxis.set_label_position('left')
ax3.set_yticks([-0.5, -0.4, -0.3, -0.2, -0.1, 0, 0.1, 0.2, 0.3])
ax3.yaxis.set_minor_locator(ticker.AutoMinorLocator(2))
ax3.xaxis.set_minor_locator(ticker.AutoMinorLocator(2))
ax3.set_xlim(0.3, 4.15)
ax3.set_ylim(-0.51, 0.11)

ax3.text(0.08, 0.065, r'$\gamma = 0$', color='#000000', transform=ax3.transAxes, fontSize=fsLabel, ha='left', bbox=dict(boxstyle='round', ec=(219/255., 202/255., 175/255.), fc=(244/255., 232/255., 211/255.)))
ax3.text(0.825, 0.985, r'$U_c/t = 3.85$', transform=ax3.transAxes, fontsize=16, verticalalignment='top', color='#888888', rotation=90)
ax3.text(0.23, 0.16, r'40% reduced $v_0$', color='#cc2909', transform=ax3.transAxes, fontSize=16, ha='left')
ax3.text(0.06, 0.9, r'${\rm Tang\; {\it et\; al.}}$', color='#000000', transform=ax3.transAxes, fontSize=21, ha='left')
ax3.annotate("", xy=(3.5, -0.4), xytext=(3., -0.4), arrowprops=dict(arrowstyle="->", linewidth=1.2, color='#cc2909'))

plt.axhline(0, color='#dddddd', zorder=-2, linewidth=0.5)
plt.axvline(3.85, linestyle='--', color='#888888', linewidth=1.0)

# The published version of this plot includes data in panel C, which was
# provided by Tang et al.; The data for was provided in pdf form only, 
# which we merged with the pdf generated from this script.

plt.subplots_adjust(wspace=8)
plt.tight_layout()
plt.savefig('Fig2.pdf')
